import ray
from pydantic import BaseModel

# 确保 Ray 已初始化
ray.init(
    address="ray://10.130.9.3:10001",
    runtime_env={
        "pip": ["pyNastran"],
        "setup_timeout_seconds": 1000,
    }
)
#ray.init()

class CustomBaseModel(BaseModel):
    def __reduce__(self):
        return (self.__class__, (), self.dict())

    @classmethod
    def __new__(cls, *args, **kwargs):
        if len(args) == 1 and not kwargs and isinstance(args[0], dict):
            # This is likely being called during deserialization
            return cls.parse_obj(args[0])
        return super().__new__(cls)

    class Config:
        arbitrary_types_allowed = True


class Parameter(CustomBaseModel):
    sample_num: int
    sample_float: float


class KeyResult:
    y: float
    rrr: str

@ray.remote
def stress_strain():
    #return 333
    res = KeyResult(y=3.14, rrr="rrr")
    return res

# 主程序
if __name__ == "__main__":
    # 创建一个 Parameter 实例
    param = Parameter(sample_num=10, sample_float=1.23)

    # 使用 Ray remote 调用 stress_strain 函数
    result_ref = stress_strain.remote()

    # 获取结果
    result = ray.get(result_ref)

    print(f"Result: y={result.y}, rrr={result.rrr}")

# 关闭 Ray
ray.shutdown()